{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "TWLmmaQpX-i1"
      },
      "source": [
        "# TensorFlow NumPy: Keras and Distribution Strategy"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "fmGBjt1arUk7"
      },
      "source": [
        "## Overview\n",
        "\n",
        "TensorFlow Numpy provides an implementation of a subset of NumPy API on top of TensorFlow backend. Please see [TF NumPy API documentation](https://www.tensorflow.org/api_docs/python/tf/experimental/numpy) and \n",
        " [TensorFlow NumPy Guide](https://www.tensorflow.org/guide/tf_numpy).\n",
        "\n",
        "This document shows how TensorFlow NumPy interoperates with TensorFlow's high level APIs like DistributionStrategky and Keras."
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "eAf_CAIerkPZ"
      },
      "source": [
        "## Setup"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "OG0u3eVdSOAk"
      },
      "outputs": [],
      "source": [
        "!pip install --quiet --upgrade tf-nightly"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "YjQUVUd3X325"
      },
      "outputs": [],
      "source": [
        "import tensorflow as tf\n",
        "import tensorflow.experimental.numpy as tnp\n",
        "\n",
        "# Creates 3 logical GPU devices for demonstrating distribution.\n",
        "gpu_device = tf.config.list_physical_devices(\"GPU\")[0]\n",
        "tf.config.set_logical_device_configuration(\n",
        "    gpu_device, [tf.config.LogicalDeviceConfiguration(128)] * 3)\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "UTZPYMaPr_oU"
      },
      "source": [
        "## TF NumPy and Keras\n",
        "\n",
        "TF NumPy can be used to create custom Keras layers. These layers interoperate with and behave like regular Keras layers. Here are some things to note to understand how these layers work.\n",
        "\n",
        "- Existing Keras layers can be invoked with ND Array inputs, in addition to other input types like `tf.Tensor`, `np.ndarray`, python literals, etc. All these types will be internally convert to a `tf.Tensor` before the layer's `call` method is invoked\n",
        "- Existing Keras layers will continue to output `tf.Tensor` values. Custom layers could output ND Array or `tf.Tensor`. \n",
        "- Custom and existing Keras layers should be freely composable.\n",
        "\n",
        "Checkout the examples below that demonstrate the above.\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "gsZLC4eEsm8P"
      },
      "source": [
        "### ND Array inputs\n",
        "\n",
        "Create and call an existing Keras layers with ND Array inputs. Note that the layer outputs a `tf.Tensor`."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "CTiylo_UrxW7"
      },
      "outputs": [],
      "source": [
        "dense_layer = tf.keras.layers.Dense(5)\n",
        "inputs = tnp.random.randn(2, 3).astype(tnp.float32)\n",
        "outputs = dense_layer(inputs)\n",
        "print(\"Shape:\", outputs.shape)\n",
        "print(\"Class:\", outputs.__class__)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "vltJnASzXJNq"
      },
      "source": [
        "### Custom Keras Layer\n",
        "\n",
        "Create a new Keras layer as below using TensorFlow NumPy methods.  Note that the layer's call method receives a `tf.tensor` value as input. It can convert to `ndarray` using `tnp.asarray`. However this conversion may not be needed since TF NumPy APIs can handle `tf.Tensor` inputs."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "0i7lOWJwsVMy"
      },
      "outputs": [],
      "source": [
        "class ProjectionLayer(tf.keras.layers.Layer):\n",
        "  \"\"\"Linear projection layer using TF NumPy.\"\"\"\n",
        "\n",
        "  def __init__(self, units):\n",
        "    super(ProjectionLayer, self).__init__()\n",
        "    self._units = units\n",
        "\n",
        "  def build(self, input_shape):\n",
        "    stddev = tnp.sqrt(self._units).astype(tnp.float32)\n",
        "    initial_value = tnp.random.randn(input_shape[1], self._units).astype(\n",
        "        tnp.float32) / stddev\n",
        "    # Note that TF NumPy can interoperate with tf.Variable.\n",
        "    self.w = tf.Variable(initial_value, trainable=True)\n",
        "\n",
        "  def call(self, inputs):\n",
        "    return tnp.matmul(inputs, self.w)\n",
        "\n",
        "# Call with ndarray inputs\n",
        "layer = ProjectionLayer(2)\n",
        "tnp_inputs = tnp.random.randn(2, 4).astype(tnp.float32)\n",
        "print(\"output:\", layer(tnp_inputs))\n",
        "\n",
        "# Call with tf.Tensor inputs\n",
        "tf_inputs = tf.random.uniform([2, 4])\n",
        "print(\"\\noutput: \", layer(tf_inputs))"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "UExEbq1EENLB"
      },
      "source": [
        "### Composing layers\n",
        "\n",
        "Next create a Keras model by composing the `ProjectionLayer` defined above with a `Dense` layer."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "qbTkqFgDDXaw"
      },
      "outputs": [],
      "source": [
        "batch_size = 3\n",
        "units = 5\n",
        "model = tf.keras.Sequential([tf.keras.layers.Dense(units),\n",
        "                             ProjectionLayer(2)])\n",
        "\n",
        "print(\"Calling with ND Array inputs\")\n",
        "tnp_inputs = tnp.random.randn(batch_size, units).astype(tnp.float32)\n",
        "output = model.call(tnp_inputs)\n",
        "print(\"Output shape %s.\\nOutput class: %s\\n\" % (output.shape, output.__class__))\n",
        "\n",
        "print(\"Calling with tensor inputs\")\n",
        "tf_inputs = tf.convert_to_tensor(tnp_inputs)\n",
        "output = model.call(tf_inputs)\n",
        "print(\"Output shape %s.\\nOutput class: %s\" % (output.shape, output.__class__))\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "QeooMJZdYbXq"
      },
      "source": [
        "## Distributed Strategy: tf.distribution\n",
        "\n",
        "[TensorFlow NumPy Guide](https://colab.sandbox.google.com/drive/15AshdHLS_xTMohWDleTiAgyPdRt6JQJJ#scrollTo=s2enCDi_FvCR) shows how `tf.device` API can be used to place individual operations on specific devices. Note that this works for remote devices as well.\n",
        "\n",
        "\n",
        "TensorFlow also has higher level distribution APIs that make it easy to replicate computation across devices. \n",
        "Here we will show how to place TensorFlow NumPy code in a Distribution Strategy context to easily perform replicated computation.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "tOTNvkTxZ-ok"
      },
      "outputs": [],
      "source": [
        "# Initialize the strategy\n",
        "gpus = tf.config.list_logical_devices(\"GPU\")\n",
        "print(\"Using following GPUs\", gpus)\n",
        "\n",
        "strategy = tf.distribute.MirroredStrategy(gpus)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "Zlmeo8i7Euq0"
      },
      "source": [
        "### Simple replication example\n",
        "\n",
        "First try running a simple NumPy function in `strategy` context."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "u3ZLh3_ZB8mk"
      },
      "outputs": [],
      "source": [
        "@tf.function\n",
        "def replica_fn():\n",
        "  replica_id = tf.distribute.get_replica_context().replica_id_in_sync_group\n",
        "  print(\"Running on device %s\" % replica_id.device)\n",
        "  return tnp.asarray(replica_id) * 5\n",
        "\n",
        "print(strategy.run(replica_fn).values)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {
        "colab_type": "text",
        "id": "UyyZBpLyE9LG"
      },
      "source": [
        "### Replicated model execution\n",
        "\n",
        "Next run the model defined earlier under `strategy` scope."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "colab": {},
        "colab_type": "code",
        "id": "6VeBFzTCCbZk"
      },
      "outputs": [],
      "source": [
        "# Test running the model in a distributed setting.\n",
        "model = tf.keras.Sequential([tf.keras.layers.Dense(units), ProjectionLayer(2)])\n",
        "\n",
        "@tf.function\n",
        "def model_replica_fn():\n",
        "  inputs = tnp.random.randn(batch_size, units).astype(tnp.float32)\n",
        "  return model.call(inputs)\n",
        "\n",
        "print(\"Outputs:\\n\", strategy.run(model_replica_fn).values)"
      ]
    }
  ],
  "metadata": {
    "accelerator": "GPU",
    "colab": {
      "collapsed_sections": [],
      "name": "TensorFlow NumPy: Keras and Distribution Strategy",
      "private_outputs": true,
      "provenance": [],
      "toc_visible": true
    },
    "kernelspec": {
      "display_name": "Python 3",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}
